module process_mod

	use mpi
	use const_mod
	use mesh_mod

	implicit none

	private
	public process_init

	type process_neighbor_type
		integer :: id       = MPI_COMM_NULL
		integer :: cart_id  = MPI_COMM_NULL
		integer :: orient   = 0

	end type process_neighbor_type

	type zonal_circle_type
		integer :: group = MPI_GROUP_NULL
		integer :: comm  = MPI_COMM_NULL
		integer :: np    = 0
		integer :: id    = MPI_PROC_NULL
		integer :: west_ngb_id = MPI_PROC_NULL
		integer :: east_ngb_id = MPI_PROC_NULL
	contains
		procedure :: init  => zonal_circle_init
		procedure :: clear => zonal_circle_clear
		final :: zonal_circle_final
	end type zonal_circle_type

	type process_type
		integer :: comm           = MPI_COMM_NULL
		integer :: cart_comm      = MPI_COMM_NULL
		integer :: group          = MPI_COMM_NULL
		integer :: cart_group     = MPI_COMM_NULL
		integer :: cart_dims(2)   = 0
		integer :: cart_coords(2) = 0
		integer :: id             = MPI_PROC_NULL ! MPI process ID 
		integer :: cart_id        = MPI_PROC_NULL ! MPI process ID in cart_comm
		integer np
		integer idom
		integer num_lon
		integer num_lat
		integer lon_ibeg
		integer lon_iend
		integer lat_ibeg
		integer lat_iend
		logical :: at_south_pole  = .false.
		logical :: at_north_pole  = .false.
		type(zonal_circle_type) zonal_circle
		type(process_neighbor_type), allocatable :: ngb(:)
	end type process_type

	type(process_type) proc

contains

	subroutine process_init()

		integer ierr

		call MPI_INIT(ierr)
		proc%comm = MPI_COMM_WORLD

		call MPI_COMM_GROUP(proc%comm, proc%group, ierr)
		call MPI_COMM_SIZE(proc%comm, proc%np, ierr)
		call MPI_COMM_RANK(proc%comm, proc%id, ierr)

		call setup_mpi_simple()
		call decompose_domains()
		call setup_zonal_comm()

	end subroutine process_init


	subroutine setup_mpi_simple()

		integer ierr, tmp_comm
		logical periods(2)

		proc%cart_dims = [1, proc%np]
		proc%idom = 1
		periods = [.true., .false.]

		call MPI_COMM_SPLIT(proc%comm, proc%idom, proc%id, tmp_comm, ierr)
		call MPI_CART_CREATE(tmp_comm, 2, proc%cart_dims, periods, .true., proc%cart_comm, ierr)
		call MPI_COMM_GROUP(proc%cart_comm, proc%cart_group, ierr)
		call MPI_COMM_FREE(tmp_comm, ierr)
		call MPI_COMM_RANK(proc%cart_comm, proc%cart_id, ierr)
		call MPI_CART_COORDS(proc%cart_comm, proc%cart_id, 2, proc%cart_coords, ierr)

	end subroutine setup_mpi_simple

	subroutine decompose_domains()
	   
	  integer ierr, tmp_id(1), i, j

	  if (allocated(proc%ngb)) deallocate(proc%ngb)
	  allocate(proc%ngb(4))
	  call MPI_CART_SHIFT(proc%cart_comm, 0, 1, proc%ngb(west )%cart_id, proc%ngb(east )%cart_id, ierr)
	  call MPI_CART_SHIFT(proc%cart_comm, 1, 1, proc%ngb(south)%cart_id, proc%ngb(north)%cart_id, ierr)

	  ! Translate Cartesian ID of neighbors to global ID.
	  do i = 1, size(proc%ngb)
	  	if (proc%ngb(i)%id == MPI_PROC_NULL) then
	  		call MPI_GROUP_TRANSLATE_RANKS(proc%cart_group, 1, [proc%ngb(i)%cart_id], proc%group, tmp_id, ierr)
	  		proc%ngb(i)%id = tmp_id(1)
	  	end if
	  end do

	  ! Handle processes at poles.
	  if (proc%ngb(south)%id == MPI_PROC_NULL) then
	  	i = proc%id + proc%cart_dims(1) / 2 * proc%cart_dims(2)
	  	if (i >= proc%np) i = i - proc%np
	  	proc%ngb(south)%id = i
	  	proc%at_south_pole = .true.
	  end if
	  if (proc%ngb(north)%id == MPI_COMM_NULL) then
	  	i = proc%id + proc%cart_dims(1) / 2 * proc%cart_dims(2)
	  	if (i >= proc%np) i = i - proc%np
	  	proc%ngb(north)%id = i
	  	proc%at_north_pole = .true.
	  end if
	  
	  ! set initial values for num_lon, num_lat, lon_ibeg, lat_ibeg
	  proc%num_lon = mesh%num_full_lon
	  proc%num_lat = mesh%num_full_lat

	  call round_robin(proc%cart_dims(1), proc%cart_coords(1), proc%num_lon, proc%lon_ibeg, proc%lon_iend)
	  call round_robin(proc%cart_dims(2), proc%cart_coords(2), proc%num_lat, proc%lat_ibeg, proc%lat_iend)

	end subroutine decompose_domains

	subroutine setup_zonal_comm()

		if (proc%idom == 1) then ! Only root domain has polar region
			call proc%zonal_circle%init()
		end if

	end subroutine setup_zonal_comm

	subroutine zonal_circle_init(this)

		class(zonal_circle_type), intent(inout) :: this

	end subroutine zonal_circle_init

	subroutine zonal_circle_clear(this)

		class(zonal_circle_type), intent(inout) :: this

	end subroutine zonal_circle_clear

	subroutine zonal_circle_final(this)

		type(zonal_circle_type), intent(inout) :: this

	end subroutine zonal_circle_final

	subroutine round_robin(dim, coord, num, ibeg, iend)

		integer, intent(in) :: dim
		integer, intent(in) :: coord
		integer, intent(inout) :: num
		integer, intent(out) :: ibeg
		integer, intent(out) :: iend

		integer res_num, tmp_num, i

    res_num = mod(num, dim)
    ibeg = 1
    do i = 0, coord - 1
      if (res_num /= 0 .and. i < res_num) then
        tmp_num = num / dim + 1
      else
        tmp_num = num / dim
      end if
      ibeg = ibeg + tmp_num
    end do
    if (res_num /= 0 .and. coord < res_num) then
      num = num / dim + 1
    else
      num = num / dim
    end if
    iend = ibeg + num - 1

	end subroutine round_robin

end module process_mod